import subprocess
import time
import sys
import os
from datetime import datetime

def get_device_id():
    output = subprocess.check_output(['hdc', 'list', 'targets']).decode('utf-8')
    lines = output.strip().split('\n')
    if lines:
        return lines[0]  # 假设第一个设备ID是我们要的
    return None

def restart_device(device_id):
    subprocess.Popen(['hdc', '-t', device_id, 'target', 'boot'])

def mount_target(device_id):
    subprocess.run(['hdc', '-t', device_id, 'target', 'mount'])

def snapshot_display(device_id):
    subprocess.run(['hdc', '-t', device_id, 'shell', 'snapshot_display -f /data/0.jpeg'])

def copy_snapshot_from_device(device_id):
    subprocess.run(['hdc', '-t', device_id, 'file', 'recv', '/data/0.jpeg'])
    
def is_reboot_successful(filename):
    current_dir = os.getcwd()
    file_path = os.path.join(current_dir, filename)       
    return os.path.exists(file_path)

def write_reboot_record(reboot_no,reboot_time, success):
    with open('reboot_record.txt', 'a') as f:
        f.write(f"{reboot_no} - {reboot_time} - {'Success' if success else 'Failure'}\n")

def delete_snapshot(filename):
    current_dir = os.getcwd()
    file_path = os.path.join(current_dir, filename)    
    if os.path.exists(file_path):
        os.remove(file_path)

def delete_snapshot_on_device(device_id):
    subprocess.run(['hdc', '-t', device_id, 'shell', 'rm -f /data/0.jpeg'])        

def main(run_times):
    device_id = get_device_id()
    if not device_id:
        print("No device found.")
        sys.stdout.flush()
        return

    for i in range(int(run_times)):
        print(f"Restarting device {i+1}/{run_times}")
        sys.stdout.flush()
        reboot_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        reboot_no = i+1
        restart_device(device_id)
        time.sleep(60)  # 睡眠1分钟

        mount_target(device_id)
        snapshot_display(device_id)
        copy_snapshot_from_device(device_id)

        if is_reboot_successful('0.jpeg'):
            print("Restart successful.")
            sys.stdout.flush()
            write_reboot_record(reboot_no,reboot_time, True)
            delete_snapshot_on_device(device_id)
            delete_snapshot('0.jpeg')
            time.sleep(60)  # 睡眠1分钟            
        else:
            print("************************ 重启失败！ *********************")
            sys.stdout.flush()
            write_reboot_record(reboot_no,reboot_time, False)
            break

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print("Usage: python restart_device.py <Run_Times>")
        sys.stdout.flush()
    else:
        main(sys.argv[1])
